import genesis as gs
import numpy as np
import torch

gs.init(backend=gs.gpu)

viewer_options = gs.options.ViewerOptions(
    camera_pos=(0, -3.5, 2.5),
    camera_lookat=(0.0, 0.0, 0.5),
    camera_fov=40,
    max_FPS=30,
)

scene = gs.Scene(
    viewer_options=viewer_options,
    sim_options=gs.options.SimOptions(
        dt=0.01,
        substeps=4, # 增加子步骤提高抓取稳定性
        gravity=(0, 0, -10),
    ),
    show_viewer=True,
    renderer=gs.renderers.Rasterizer(),
)

ground = scene.add_entity(
    gs.morphs.Plane()
)

franka = scene.add_entity(
    gs.morphs.MJCF(
        file="xml/franka_emika_panda/panda.xml",
        pos = (0, 0, 0),
        euler = np.random.rand(3) * 2 * np.pi,
        scale = 1.0,          
    ),
)

sphere_radius = 0.03
sphere_pos = np.array([0.65, 0.0, sphere_radius])

target_sphere = scene.add_entity(
    gs.morphs.Sphere(
        radius=sphere_radius,
        pos=sphere_pos
    ),
)

scene.build()

joint_names = [
    "joint1",
    "joint2",
    "joint3",
    "joint4",
    "joint5",
    "joint6",
    "joint7",
    "finger_joint1",
    "finger_joint2",
]

dofs_idx = [franka.get_joint(name).dof_idx_local for name in joint_names]



############ 可选：设置控制增益 ############
# 设置位置增益
franka.set_dofs_kp(
    kp             = np.array([4500, 4500, 3500, 3500, 2000, 2000, 2000, 100, 100]),
    dofs_idx_local = dofs_idx,
)
# 设置速度增益
franka.set_dofs_kv(
    kv             = np.array([450, 450, 350, 350, 200, 200, 200, 10, 10]),
    dofs_idx_local = dofs_idx,
)
# 设置安全的力范围
franka.set_dofs_force_range(
    lower          = np.array([-87, -87, -87, -87, -12, -12, -12, -100, -100]),
    upper          = np.array([ 87,  87,  87,  87,  12,  12,  12,  100,  100]),
    dofs_idx_local = dofs_idx,
)

hand = franka.get_link("hand")

franka.control_dofs_position(
    torch.rand(9, device=gs.device),
    dofs_idx,
)

# 先运动到小球的上方0.05处
height = 0.2
firset_step_pos = sphere_pos + np.array([0, 0, height])
qpos = franka.inverse_kinematics(
    link = hand,
    pos = firset_step_pos,  # 这里手指不能直接给到球的位置，否则会发生碰撞，倒置规划失败
    quat = np.array([0,0,1,0]),
)
print(f"Initial qpos: {qpos}")
print(f"Initial qpos: {qpos}")

# 打开夹爪
qpos[-2:] = 1


path = franka.plan_path(
    qpos_goal = qpos,
    num_waypoints = 200,
)

if path is None:
    print("No path found")
    exit()

for waypoint in path:
    franka.control_dofs_position(waypoint)
    scene.step()

for i in range(50):
    scene.step()

# 到达
height = 0.1
second_step_pos = sphere_pos + np.array([0, 0, height])
qpos = franka.inverse_kinematics(
    link = hand,
    pos  = second_step_pos, # 求出当前的pos
    quat = np.array([0, 0, 1, 0]),
)

body_dofs = np.arange(0,7)
hand_dofs = np.arange(7,9)
franka.control_dofs_position(qpos[0:-2], body_dofs)
for i in range(50):
    scene.step()

franka.control_dofs_force(np.array([-20, -20]), hand_dofs)

for i in range(50):
    scene.step()

# 抓取完成
target_pos = np.array([0.65, 0, 0.3])
qpos = franka.inverse_kinematics(
    link = hand,
    pos  = target_pos, # 求出当前的pos
    quat = np.array([0, 0, 1, 0]),
)
# 抬起
franka.control_dofs_position(qpos, dofs_idx)
for i in range(50):
    scene.step()

franka.control_dofs_force(np.array([0,0]), hand_dofs)
for i in range(50):
    scene.step()

lift_pose = qpos.clone()

# 移动到对称位置
target_pos = np.array([-0.65, 0.1, 0.3])
qpos = franka.inverse_kinematics(
    link = hand,
    pos  = target_pos, # 求出当前的pos
    quat = np.array([0, 0, 1, 0]),
    init_qpos = qpos,
)
path = franka.plan_path(
    qpos_goal = qpos,
    qpos_start = lift_pose,
    num_waypoints = 200,
)
for waypoint in path:
    franka.control_dofs_position(waypoint[0:-2], body_dofs)
    scene.step()

for i in range(10000):
    scene.step()